#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
历史数据导入SQL数据库脚本

用于从CSV文件中导入数据到SQL数据库
"""

import argparse
import os
import sys
import pandas as pd
import glob
import logging
from sqlalchemy import create_engine
from sqlalchemy import text
# 配置日志
logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s',
    level=logging.INFO
)
logger = logging.getLogger('import_historical_data')

def create_table(engine, table_name):
    """
    创建数据库表结构
    
    参数:
        engine: SQLAlchemy引擎
        table_name: 表名
    """
    stat = text(f"""
        CREATE TABLE IF NOT EXISTS {table_name} (
            id INT AUTO_INCREMENT PRIMARY KEY,
            date VARCHAR(8) NOT NULL,
            hour INT NOT NULL,
            type VARCHAR(50) NOT NULL,
            site_id VARCHAR(50) NOT NULL,
            value FLOAT,
            UNIQUE KEY unique_record (date, hour, type, site_id)
        )
        """)
    with engine.connect() as conn:
        conn.execute(stat)
    logger.info(f"已创建表: {table_name}")

def import_csv_to_sql(csv_file, engine, table_name):
    """
    从CSV文件导入数据到SQL数据库
    
    参数:
        csv_file: CSV文件路径
        engine: SQLAlchemy引擎
        table_name: 表名
    
    返回:
        导入的记录数量
    """
    if not os.path.exists(csv_file):
        logger.error(f"文件不存在: {csv_file}")
        return 0
    
    try:
        # 读取CSV文件
        logger.info(f"正在读取文件: {csv_file}")
        df = pd.read_csv(csv_file)
        logger.info(f"文件读取完成: {os.path.basename(csv_file)}, {len(df)} 条记录")
        
        # 获取所有列名
        all_columns = df.columns.tolist()
        base_columns = ['date', 'hour', 'type']
        site_columns = [col for col in all_columns if col not in base_columns]
        
        # 将宽表转换为长表
        df_melted = pd.melt(df, 
                          id_vars=base_columns,
                          value_vars=site_columns,
                          var_name='site_id',
                          value_name='value')
        
        # 导入到数据库
        df_melted.to_sql(table_name, 
                        engine, 
                        if_exists='append', 
                        index=False,
                        chunksize=1000)
        
        logger.info(f"成功导入 {len(df_melted)} 条记录到表 {table_name}")
        return len(df_melted)
    
    except Exception as e:
        logger.error(f"导入文件 {csv_file} 失败: {str(e)}")
        return 0

def import_directory_to_sql(csv_dir, engine, table_name):
    """
    从目录导入所有CSV文件到SQL数据库
    
    参数:
        csv_dir: CSV文件目录
        engine: SQLAlchemy引擎
        table_name: 表名
    
    返回:
        处理的文件数量和总记录数
    """
    if not os.path.exists(csv_dir):
        logger.error(f"目录不存在: {csv_dir}")
        return 0, 0
    
    # 获取所有CSV文件
    csv_files = glob.glob(os.path.join(csv_dir, "*.csv"))
    if not csv_files:
        logger.error(f"未找到CSV文件: {csv_dir}")
        return 0, 0
    
    logger.info(f"找到 {len(csv_files)} 个CSV文件")
    
    # 创建表
    create_table(engine, table_name)
    
    processed_files = 0
    total_records = 0
    for csv_file in csv_files:
        records = import_csv_to_sql(csv_file, engine, table_name)
        if records > 0:
            processed_files += 1
            total_records += records
    
    logger.info(f"总共处理了 {processed_files} 个文件, 导入 {total_records} 条记录")
    return processed_files, total_records

def main():
    # 解析命令行参数
    parser = argparse.ArgumentParser(description='从CSV文件导入数据到SQL数据库')
    parser.add_argument('--file', type=str, help='单个CSV文件路径')
    parser.add_argument('--dir', type=str, help='CSV文件目录')
    parser.add_argument('--db-url', type=str, required=False, 
                       help='数据库连接URL，例如: mysql://user:password@localhost/dbname')
    parser.add_argument('--table', type=str, default='air_quality', 
                       help='表名，默认为air_quality')
    args = parser.parse_args()
    
    # 创建数据库引擎
    if not args.db_url:
        from config.db_config import DB_URL
        engine = create_engine(DB_URL)
    else:
        engine = create_engine(args.db_url)
    
    if args.file:
        # 处理单个文件
        import_csv_to_sql(args.file, engine, args.table)
    elif args.dir:
        # 处理目录中的所有文件
        import_directory_to_sql(args.dir, engine, args.table)
    else:
        logger.error("必须指定--file或--dir参数")
        parser.print_help()
        return 1
    
    return 0

if __name__ == "__main__":
    sys.exit(main())